import os
import shutil
import argparse
import sys


def backup_large_files(source_dir, backup_dir, size_threshold_kb=32):
    """
    备份超过指定大小的文件并删除原文件
    :param source_dir: 源目录
    :param backup_dir: 备份目录
    :param size_threshold_kb: 大小阈值(KB)
    """
    if not os.path.exists(backup_dir):
        os.makedirs(backup_dir)

    size_threshold = size_threshold_kb * 1024  # 转换为字节
    backed_up_files = []

    for root, _, files in os.walk(source_dir):
        for file in files:
            file_path = os.path.join(root, file)
            if not os.path.isfile(file_path):
                continue

            try:
                file_size = os.path.getsize(file_path)
                is_triton_jit = False
                with open(file_path, 'r', encoding='utf-8') as file:
                    for _, line in enumerate(file, 1):
                        if '@triton.jit' in line:
                            is_triton_jit = True
                            break
            except OSError:
                continue

            if file_size > size_threshold or is_triton_jit:
                # 计算备份路径，保持目录结构
                rel_path = os.path.relpath(file_path, source_dir)
                backup_path = os.path.join(backup_dir, rel_path)

                # 创建必要的目录
                os.makedirs(os.path.dirname(backup_path), exist_ok=True)

                # 备份文件
                shutil.copy2(file_path, backup_path)
                backed_up_files.append((file_path, backup_path))

                # 删除原文件
                os.remove(file_path)
                print(f"Backed up and removed: {file_path} -> {backup_path}")

    # 保存备份记录
    if backed_up_files:
        with open(os.path.join(backup_dir, 'backup_manifest.txt'), 'w') as f:
            for original, backup in backed_up_files:
                f.write(f"{original}|{backup}\n")

    return backed_up_files



def restore_backup(backup_dir, target_dir):
    """
    从备份还原文件
    :param backup_dir: 备份目录
    :param target_dir: 目标目录
    """
    manifest_file = os.path.join(backup_dir, 'backup_manifest.txt')

    if not os.path.exists(manifest_file):
        print("Error: No backup manifest found. Cannot restore.")
        return

    restored_count = 0
    with open(manifest_file, 'r') as f:
        for line in f:
            try:
                original_path, backup_path = line.strip().split('|')

                # 确保目标目录存在
                os.makedirs(os.path.dirname(original_path), exist_ok=True)

                # 检查备份文件是否存在
                if not os.path.exists(backup_path):
                    print(f"Warning: Backup file not found: {backup_path}")
                    continue

                # 还原文件
                shutil.copy2(backup_path, original_path)
                restored_count += 1
                print(f"Restored: {backup_path} -> {original_path}")
            except ValueError:
                print(f"Warning: Invalid line in manifest: {line.strip()}")

    print(f"Restore completed. {restored_count} files restored.")


def main():
    parser = argparse.ArgumentParser(description="文件备份、加密和还原工具")
    subparsers = parser.add_subparsers(dest='command', required=True)

    # 备份命令
    backup_parser = subparsers.add_parser('backup', help='备份大文件并删除原文件')
    backup_parser.add_argument('source_dir', help='源目录')
    backup_parser.add_argument('backup_dir', help='备份目录')
    backup_parser.add_argument('--size', type=int, default=32,
                               help='大小阈值(KB)，默认32KB')

    # 还原命令
    restore_parser = subparsers.add_parser('restore', help='从备份还原文件')
    restore_parser.add_argument('backup_dir', help='备份目录')
    restore_parser.add_argument('target_dir', help='还原目标目录')

    args = parser.parse_args()

    try:
        if args.command == 'backup':
            backup_large_files(args.source_dir, args.backup_dir, args.size)
            print("备份完成")
        elif args.command == 'restore':
            restore_backup(args.backup_dir, args.target_dir)
            print("还原完成")
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()